package com.lst.selfaiagent.advisors;

import org.springframework.ai.chat.client.advisor.api.*;
import reactor.core.publisher.Flux;

import java.util.Arrays;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;


public class VerificationAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {


    // 违禁词列表，实际应用中可以从配置文件或数据库加载
    private static final Set<String> BANNED_WORDS = new HashSet<>(Arrays.asList(
            "暴力", "色情", "赌博", "毒品", "违法", "漏骨", "擦边"
    ));

    private AdvisedRequest before(AdvisedRequest advisedRequest) {
        String originalText = advisedRequest.userText();

        // 校验违禁词
        String filteredText = checkAndFilterBannedWords(originalText);

        // 优化提示词
        String optimizedText = optimizePrompt(filteredText);

        return AdvisedRequest.from(advisedRequest)
                .userText(optimizedText)
                .build();
    }


    @Override
    public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {

        advisedRequest = this.before(advisedRequest);
        AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);
//        advisedResponse = this.after(advisedResponse);
        return advisedResponse;
    }


    @Override
    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
        advisedRequest = this.before(advisedRequest);
        return chain.nextAroundStream(advisedRequest);

    }

    @Override
    public String getName() {
        return "违禁词校验顾问";
    }

    @Override
    public int getOrder() {
        return 0;
    }

    /**
     * 检查并过滤违禁词
     * @param text 原始文本
     * @return 过滤后的文本
     * @throws SecurityException 如果包含违禁词则抛出异常
     */
    private String checkAndFilterBannedWords(String text) {
        if (text == null || text.isEmpty()) {
            return text;
        }

        // 检查是否包含违禁词
        for (String bannedWord : BANNED_WORDS) {
            if (text.contains(bannedWord)) {
                // 可以选择抛出异常或替换违禁词
                // 抛出异常方式:
                // throw new SecurityException("输入内容包含违禁词: " + bannedWord);

                // 替换方式(这里用**替换违禁词):
                text = text.replace(bannedWord, "*".repeat(bannedWord.length()));
            }
        }

        return text;
    }

    /**
     * 优化提示词
     * @param originalText 原始文本
     * @return 优化后的文本
     */
    private String optimizePrompt(String originalText) {
        if (originalText == null || originalText.isEmpty()) {
            return originalText;
        }

        // 基本的提示词优化逻辑
        String optimizedText = originalText.trim();

        // 移除多余的空格和换行
        optimizedText = optimizedText.replaceAll("\\s+", " ");

        // 添加清晰度和具体性提示（如果需要）
        if (!optimizedText.contains("请") && !optimizedText.contains("请")) {
            optimizedText = "请" + optimizedText;
        }

        // 确保句子以合适的标点结束
        if (!optimizedText.endsWith("？") && !optimizedText.endsWith("。") && !optimizedText.endsWith("!")) {
            if (optimizedText.contains("？") || optimizedText.contains("?")) {
                optimizedText += "？";
            } else {
                optimizedText += "。";
            }
        }

        return optimizedText;
    }
}
